import time

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torch import nn 
import numpy as np
import copy
from models.Update import DatasetSplit
from utils.save_result import save_result
from models.aggregation import Aggregation
from models.test import test_img_avg
from models.model_creator import *
from models.ResNet_Dery import *
from sim.block_meta import MODEL_BLOCKS,MODEL_ZOO,PLANES 
from collections import OrderedDict
import wandb
class LocalUpdate_FedBase(object):
    def __init__(self, args, dataset=None, idxs=None, verbose=False):
        self.args = args
        self.loss_func = nn.CrossEntropyLoss()
        self.selected_clients = []
        self.ldr_train = DataLoader(DatasetSplit(dataset, idxs,args), batch_size=self.args.local_bs, shuffle=True)
        self.verbose = verbose

    def train(self, round, net):

        net.train()
        # train and update
        if self.args.optimizer == 'sgd':
            optimizer = torch.optim.SGD(net.parameters(), lr=self.args.lr*(self.args.lr_decay**round),
                                        momentum=self.args.momentum,weight_decay=self.args.weight_decay)
        elif self.args.optimizer == 'adam':
            optimizer = torch.optim.Adam(net.parameters(), lr=self.args.lr)

        Predict_loss = 0
        for iter in range(self.args.local_ep):

            for batch_idx, (images, labels) in enumerate(self.ldr_train):
                images, labels = images.to(self.args.device), labels.to(self.args.device)
                if self.args.dataset == 'widar':
                    labels = labels.long()
                if (images.size()[0] <= 1):
                    continue
                net.zero_grad()
                log_probs = net(images)['output']
                loss = self.loss_func(log_probs, labels)
                loss.backward()
                optimizer.step()

                Predict_loss += loss.item()

        if self.verbose:
            info = '\nUser predict Loss={:.4f}'.format(Predict_loss / (self.args.local_ep * len(self.ldr_train)))
            print(info)

        return net.state_dict()

def FedBase(args,net_zoo, dataset_train, dataset_test, dict_users):

    num_model = len(net_zoo)
    # num_group = 2
    block_list = []
    block_list = split_block(copy.deepcopy(net_zoo),args.num_block)
    model_list = []
    concat_para = []
    for i in range(num_model):
        model_list.append([])
        concat_para.append([])
        for j in range(num_model):
            choose = []
            block_choice = []
            for k in range(args.num_block):
                if k < 2:
                    choose.append(i)
                else:
                    choose.append(j)
                block_choice.append(copy.deepcopy(block_list[choose[k]][k]))
            print(choose)
            model_list[i].append(CombinedModel(block_list=block_choice,choose = choose,planes=PLANES).to(args.device))
            concat_para[i].append(None)
    # training
    acc = []
    acc.append([])
    ###### choose model
    net_glob = copy.deepcopy(model_list[args.num_id1][args.num_id2] )
    net_glob.train()
    for iter in range(args.epochs):
    # for iter in range(2):
        start_time = time.time()
        print('*'*80)
        print('Round {:3d}'.format(iter))


        w_locals = []
        lens = []
        m = max(int(args.frac * args.num_users), 1)
        # m = num_model * num_group
        # idxs_users = np.random.choice(range(args.num_users), m, replace=False)
        m1 = int(m * args.client1_frac)
        m2 = int(m * args.client2_frac)
        m3 = int(m * args.client3_frac)
        m4 = m1 + m2 + m3
        idxs_users = []
        if args.num_id2 == 0:
            idxs_users.append(np.random.choice(range(0,int(args.num_users)),m4,replace=False))
        elif args.num_id2 == 1:
            if m4 > int(args.num_users) - int(args.num_users * args.client1_frac):
                m4 = int(args.num_users) - int(args.num_users * args.client1_frac)
            idxs_users.append(np.random.choice(range(int(args.num_users*args.client1_frac),int(args.num_users)),m4,replace=False))
        elif args.num_id2 == 2:
            if m4 > int(args.num_users) - int(args.num_users * (args.client1_frac + args.client2_frac)):
                m4 = int(args.num_users) - int(args.num_users * (args.client1_frac + args.client2_frac))
            idxs_users.append(np.random.choice(range(int(args.num_users*(args.client1_frac+args.client2_frac)),int(args.num_users)),m4,replace=False))
        for num,client_list in enumerate(idxs_users):
            for index,idx in enumerate(client_list):
                local = LocalUpdate_FedBase(args=args, dataset=dataset_train, idxs=dict_users[idx])
                w = local.train(round=iter,net=copy.deepcopy(net_glob).to(args.device))

                w_locals.append(copy.deepcopy(w))
                lens.append(len(dict_users[idx]))
        # update global weights
        w_glob = Aggregation(w_locals, lens)
        # acc_dict = {}
        # copy weight to net_glob
        net_glob.load_state_dict(w_glob)

        end_time = time.time()
        print("Training time: ",end_time-start_time)
        acc_item = test(net_glob, dataset_test, args)
        # print(type(acc_item))
        acc[0].append(acc_item)
        # wandb.log({'accuracy_local{:d}'.format(args.num_id1*num_model+args.num_id2+1):acc_item})
        # acc_dict['accuracy'] = acc_item

        # for i in range(m):
        #     net_glob.load_state_dict(w_locals[i])
        #     test_a = test(net_glob,dataset_test,args)
        #     acc_dict['accuracy_local{}'.format(i+1)] = test_a
        # net_glob.load_state_dict(w_glob)
        # wandb.log(acc_dict)
    save_result(acc, 'test_acc', args)

def split_block(net_zoo,num_block):
    block_list = []  ## size = num_block
    for i in range(len(net_zoo)):
        block_list.append([])
        for j in range(num_block):
            block_list[i].append([])
            for layer in MODEL_BLOCKS[MODEL_ZOO[i]][j]:
                if '.' not in layer:
                    layer1 = layer
                    layer2 = None
                else:
                    layers = layer.split('.')
                    layer1 = layers[0]
                    layer2 = layers[1]
                for name,child in net_zoo[i].named_children():
                    if layer1 in name :
                        if isinstance(child,torch.nn.Sequential) :
                            block_list[i][j].append(child[int(layer2)])
                        else:
                            block_list[i][j].append(child)
    for i in range(len(net_zoo)):
        for j in range(num_block):
            block_list[i][j] = nn.Sequential(*block_list[i][j])
    return block_list

def split_block_simple(net_zoo,num_block):
    block_list = []
    for i in range(len(net_zoo)):
        block_list.append([])
        for j in range(num_block):
            block_list[i].append([])
            for name,child in net_zoo[i].named_children():
                if 'block{:d}'.format(j+1) in name:
                    block_list[i][j].append(child)
    return block_list 

def test(net_glob, dataset_test, args):
    
    # testing
    acc_test, loss_test = test_img_avg(net_glob, dataset_test, args)

    print("Testing accuracy: {:.2f}".format(acc_test))

    return acc_test.item()